// Copyright (c) 2025 Huawei Technologies Co., Ltd.
// openUBMC is licensed under Mulan PSL v2.
// You can use this software according to the terms and conditions of the Mulan PSL v2.
// You may obtain a copy of Mulan PSL v2 at:
//         http://license.coscl.org.cn/MulanPSL2
// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
// EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
// MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
// See the Mulan PSL v2 for more details.
//
// Description: VMM虚拟媒体客户端

#include <cstdlib>
#include <iostream>
#include <thread>
#include <memory>
#include <string>
#include "VmmClient.hpp"

enum class DEVICE_TYPE : uint8_t {
    CD = 0x00,
    FLOPPY
};

VmmClient::VmmClient() = default;

VmmClient::~VmmClient() {
    Disconnect();

    isConnected_ = false;
    connectState_ = ConnectState::INIT;
    filePath_.clear();

    // 销毁 ImageDevice 和 ScsiProcessor
    scsiProcessor_.reset();
    imageDevice_.reset();
}

// 注册错误处理函数
void VmmClient::SetErrorHandler(std::function<void(const std::exception&)> handler) {
    errorHandler_ = std::move(handler);
}

// 注册会话关闭处理函数
void VmmClient::SetCloseHandler(std::function<void(const uint16_t&)> handler) {
    closeHandler_ = std::move(handler);
}

// 注册日志处理函数
void VmmClient::SetLogHandler(std::function<void(const std::string&)> handler) {
    logHandler_ = std::move(handler);
}

int VmmClient::OpenVmm(const std::string& filePath, MountType type, std::string host, uint16_t port, const char* token) {
    // 防止重入
    std::lock_guard<std::mutex> lock(openVmmMutex_);

    if (isConnected_) {
        LogMessage("OpenKvm: Already connected to VMM.");
        return E_ALREADY_CONNECTED;
    }

    try {
        connectState_ = ConnectState::INIT;
        Connect(host, std::to_string(port), "/websocket/vmm");
        std::thread t([this]() { Run(); });
        std::this_thread::sleep_for(std::chrono::seconds(1));   // 稍微等待1秒保证I/O服务已启动
        
        PackData pack_data;
        // 初始化连接响应变量
        connectRsp_ = INIT_RESPONSE_CODE;
    
        Send(pack_data.ConnectVmm(token));
        connectState_ = ConnectState::CERTIFY;
        // 发送连接请求后，等待服务端的响应
        std::unique_lock<std::mutex> rspLock(connectRspMutex_);
        if (!connectRspCv_.wait_for(rspLock, std::chrono::seconds(RESPONSE_TIMEOUT),
            [this]() { return connectRsp_ != INIT_RESPONSE_CODE; })) {
            LogMessage("OpenVmm: Timeout waiting for connection response.");
            return E_TIMEOUT;
        }

        t.detach();
        
        // 返回服务端的结果
        if (connectRsp_ == E_OK) {
            isConnected_ = true;
            filePath_ = filePath;
            mountType_ = type;

            // 初始化 ImageDevice 和 ScsiProcessor
            imageDevice_ = std::make_unique<ImageDevice>();
            scsiProcessor_ = std::make_unique<ScsiProcessor>(*imageDevice_);

            // 设置 ScsiProcessor 的回调函数
            scsiProcessor_->SetSendHandler([this](const std::vector<uint8_t>& data) {
                Send(data);
            });
            scsiProcessor_->SetLogHandler([this](const std::string& message) {
                LogMessage(message);
            });
            imageDevice_->SetFilePath(type, filePath);
            LogMessage("Open VMM successfully.");
        } else {
            LogMessage("Open VMM failed.");
            Disconnect();
        }

        return connectRsp_;
    }
    catch(const std::exception& e) {
        LogMessage(std::string("Exception in OpenVmm: ") + e.what());
        OnError(e);
        return E_ERR;
    }
}

int VmmClient::CloseVmm() {
    if (!isConnected_) {
        LogMessage("VmmDisconnect: Not connected to VMM.");
        return E_NOT_CONNECTED;
    }
    Send(PackData().DisconnectVmm());

    return E_OK;
}

void VmmClient::OnClose(const websocket::close_reason& cr) {
    isConnected_ = false;
    connectState_ = ConnectState::INIT;
    filePath_.clear();

    // 销毁 ImageDevice 和 ScsiProcessor
    scsiProcessor_.reset();
    imageDevice_.reset();

    if (closeHandler_) {
        closeHandler_(cr.code);
    }
}

void VmmClient::OnError(const std::exception& error) {
    Disconnect();

    isConnected_ = false;
    connectState_ = ConnectState::INIT;
    filePath_.clear();

    // 销毁 ImageDevice 和 ScsiProcessor
    scsiProcessor_.reset();
    imageDevice_.reset();

    if (errorHandler_) {
        errorHandler_(error);
    }
}

void VmmClient::LogMessage(const std::string& message) {
    if (logHandler_) {
        logHandler_(message);
    }
}

void VmmClient::OnMessage(const std::string& message) {
    if (message.empty()) {
        LogMessage("Received empty message.");
        return;
    }

    VMMMessage vmmMsg;
    if (!vmmMsg.Unpack(message)) {
        LogMessage("Failed to unpack message.");
        return;
    }

    uint8_t opCode = vmmMsg.GetOpCode();

    try {
        switch (opCode) {
            case ACK:
                ProcessAck(vmmMsg.GetSack());
                break;
            case TRANS_UFI:
                break;
            case TRANS_SFF:
                scsiProcessor_->ProcessScsiCommand(vmmMsg);
                break;
            case DISCONNECT:
                ProcessClose(vmmMsg.GetStatus());
                break;
            case HEARTBEAT:
                Send(PackData().VmmHeartBeat());
                break;
            case CLOSE_DEV:
                ProcessClose(0);
                break;
            default:
                break;
        }
    } catch(const std::exception& e) {
        LogMessage(std::string("Exception in OnMessage: ") + e.what());
        if (errorHandler_) {
            errorHandler_(e);
        }
    }
}

void VmmClient::ProcessAck(uint8_t sack) {
    if (connectState_ == ConnectState::CERTIFY) {
        std::lock_guard<std::mutex> lock(connectRspMutex_);
        switch (sack) {
        case OPT_SUCCESS:
            LogMessage("Connect VMM successfully");
            connectRsp_ = E_OK;
            connectState_ = ConnectState::CONNECT;
            Send(PackData().CreateDevice(static_cast<uint8_t>(DEVICE_TYPE::CD)));
            break;
        case OPT_ERROR:
            LogMessage("Connect VMM failed");
            connectRsp_ = E_ERR;
            break;
        case CONNECT_EXIST:
            LogMessage("VMM has already connected");
            connectRsp_ = E_VMM_BUSY;
            break;
        case AUTH_FAIL:
            LogMessage("Auth VMM failed");
            connectRsp_ = E_AUTHORIZATION_FAILED;
            break;
        case NO_PRIV:
            LogMessage("No permission to connect VMM");
            connectRsp_ = E_INSUFFICIENT_PRIVILEGE;
            break;
        default:
            break;
        }
        connectRspCv_.notify_one(); // 通知等待的线程
        return;
    }

    if (connectState_ == ConnectState::CONNECT || connectState_ == ConnectState::ACTIVE) {
        if (sack == CREATE_DEV_SUCCESS) {
            LogMessage("Create VMM device successfully");
            connectState_ = ConnectState::ACTIVE;
        } else {
            LogMessage("Create VMM device failed, sack: " + std::to_string(sack));
        }
    }
}

void VmmClient::ProcessClose(uint8_t status) {
    isConnected_ = false;
    connectState_ = ConnectState::INIT;

    if (closeHandler_) {
        closeHandler_(status);
    }
}

